import os
import re
import argparse
from pyfaidx import Fasta
from vcf import Reader as vcfReader, Writer as vcfWriter
from vcf.parser import _Info as vcfInfo


def get_options():
    parser = argparse.ArgumentParser(description='AutoPVS1 Batch')
    parser.add_argument('--input', '-i', required=True, help='input file, VCF for SNV, BED for CNV')
    parser.add_argument('--output', '-o', required=True, help='output file')
    parser.add_argument('--genome', '-g', required=True, help='path to the input genome fasta')
    parser.add_argument('--builder', '-b', choices=('hg19', 'hg38'), required=True, help='builder version')
    parser.add_argument('--vep_cache', '-d', required=True, help='vep cached path')
    parser.add_argument('--vartype', '-t', default='snv', choices=('snv', 'cnv'), help='vartype cnv or snv')
    parser.add_argument('--process', '-p', default=os.cpu_count(), type=int, help='process number')
    # parser.add_argument('--with_chr', '-c',  action='store_true', help='vep cached path')
    return parser.parse_args()


def autopvs1_snv(in_vcf: str, vep_cache: str, builder: str, genome_fasta: str, out_vcf: str, process:int):
    from .vep import create_snv_vep_input, run_vep_snv, read_vep_out
    from .autoPVS1 import AutoPVS1
    vep_input = f'{out_vcf}.vep_input.vcf'
    vep_output = f'{out_vcf}.vep_output.txt'
    create_snv_vep_input(in_vcf=in_vcf, vep_input=vep_input)
    run_vep_snv(vep_input=vep_input, vep_output=vep_output, vep_cache=vep_cache, vep_assembly='GRCh37' if builder == 'hg19' else 'GRCh38', process=process)
    vep_out_dict = read_vep_out(vep_output)
    genome = Fasta(genome_fasta, rebuild=False)
    reader = vcfReader(filename=in_vcf, encoding='utf-8')
    with open(out_vcf, 'w') as fo:
        if reader.infos:
            reader.infos['AutoPVS1'] = vcfInfo(id='AutoPVS1', num=None, type='String', desc='AutoPVS1,format=criterion|strength_raw|strength', source=None, version=None)
        writer = vcfWriter(fo, template=reader)
        for row in reader:
            key = f'{row.CHROM}-{row.POS}-{row.REF}-{row.ALT[0]}'
            vep_outs = vep_out_dict.get(key)
            if vep_outs:
                autopvs1 = AutoPVS1(vcf_var=key, vep_outs=vep_outs, genome_version=builder, genome=genome)
                if autopvs1.pvs1:
                    row.INFO['AutoPVS1'] = f'{autopvs1.pvs1.criterion}|{autopvs1.pvs1.strength_raw.value}|{autopvs1.pvs1.strength.value}'
            writer.write_record(row)
    os.remove(vep_input)
    os.remove(vep_output)


def autopvs1_cnv(in_vcf: str, vep_cache: str, builder: str, out_vcf: str, process:int):
    from .vep import create_cnv_vep_input, run_vep_cnv, read_vep_out
    from .autoPVS1 import AutoPVS1CNV
    vep_input = f'{out_vcf}.vep_input.txt'
    vep_output = f'{out_vcf}.vep_output.txt'
    create_cnv_vep_input(in_vcf=in_vcf, vep_input=vep_input)
    run_vep_cnv(vep_input=vep_input, vep_output=vep_output, vep_cache=vep_cache, vep_assembly='GRCh37' if builder == 'hg19' else 'GRCh38', process=process)
    vep_out_dict = read_vep_out(vep_output)
    reader = vcfReader(filename=in_vcf, encoding='utf-8')
    with open(out_vcf, 'w') as fo:
        if reader.infos:
            reader.infos['AutoPVS1'] = vcfInfo(id='AutoPVS1', num=None, type='String', desc='AutoPVS1,format=criterion|strength_raw|strength', source=None, version=None)
        writer = vcfWriter(fo, template=reader)
        for row in reader:
            key = f'{row.CHROM}-{row.POS}-{row.INFO["END"]}-{row.ALT[0]}'
            vep_outs = vep_out_dict.get(key)
            if vep_outs:
                autopvs1 = AutoPVS1CNV(bed_var=key, vep_outs=vep_outs, genome_version=builder)
                if autopvs1.pvs1:
                    row.INFO['AutoPVS1'] = f'{autopvs1.pvs1.criterion}|{autopvs1.pvs1.strength_raw.value}|{autopvs1.pvs1.strength.value}'
                    writer.write_record(row)
    os.remove(vep_input)
    os.remove(vep_output)


def main():
    args = get_options()
    if args.vartype == 'snv':
        autopvs1_snv(in_vcf=args.input, vep_cache=args.vep_cache, builder=args.builder, genome_fasta=args.genome, out_vcf=args.output, process=args.process)
    elif args.vartype == 'cnv':
        autopvs1_cnv(in_vcf=args.input, vep_cache=args.vep_cache, builder=args.builder, out_vcf=args.output, process=args.process)
    else:
        raise Exception('ERROR: make sure your input is VCF for SNV or BED for CNV')


if __name__ == '__main__':
    main()
